import os
import glob
import random
import xml.etree.ElementTree as ET

config = {
    # Annotations path(Annotations 的文件夹路径)
    "Annotation": "F:\\datasets\\breast\\labelimg\\voc\\annotations",
    # JPEGImages path(JPEGImages 的文件夹路径)
    "JPEGImages": "F:\\datasets\\breast\\labelimg\\voc\\images",
}
# 划分数据集
# 数据划分比例
# (训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1
# 按照比例划分数据集
train_per = 0.8
valid_per = 0.1
test_per = 0.1
# 读取xml
data_xml_list = glob.glob(os.path.join(config['Annotation'], '*.xml'))
# 读取图片
data_jpg_list = glob.glob(os.path.join(config['JPEGImages'], '*.*'))
data_xml_list.sort()
data_jpg_list.sort()
# 生成label标签:
label = set()
for xml_path in data_xml_list:
    label = label | set([i.find('name').text for i in ET.parse(xml_path).findall('object')])
data_list = []
for i, j in zip(data_jpg_list, data_xml_list):
    data_list.append(i + " " + j)
data_xml_list = data_list
print(data_xml_list)
random.seed(666)
random.shuffle(data_xml_list)
data_length = len(data_xml_list)

train_point = int(data_length * train_per)
train_valid_point = int(data_length * (train_per + valid_per))

# 生成训练集，验证集, 测试集(8 : 1 : 1)
train_list = data_xml_list[:train_point]
valid_list = data_xml_list[train_point:train_valid_point]
test_list = data_xml_list[train_valid_point:]

# 写入文件中
ftrain = open('/home/aistudio/data/data94796/trainval.txt', 'w')
fvalid = open('/home/aistudio/data/data94796/valid.txt', 'w')
ftest = open('/home/aistudio/data/data94796/test.txt', 'w')
flabel = open('/home/aistudio/data/data94796/label_list.txt', 'w')

for i in train_list:
    ftrain.write(i + "\n")
for j in valid_list:
    fvalid.write(j + "\n")
for k in test_list:
    ftest.write(k + "\n")
for l in label:
    flabel.write(l + "\n")
ftrain.close()
fvalid.close()
ftest.close()
flabel.close()
print(
    "总数据量:{}, 训练集:{}, 验证集:{}, 测试集:{}, 标签:{}".format(len(data_xml_list), len(train_list), len(valid_list),
                                                                   len(test_list), len(label)))
print("done!")
